# 导入Ollama聊天模型
from langchain_ollama.chat_models import ChatOllama
# 导入聊天场景的提示模板
from langchain_core.prompts import ChatPromptTemplate
# 创建一个 ChatOllama 对象

llm = ChatOllama(
    base_url="http://127.0.0.1:11434/",
    model="qwen2.5:14b",
    temperature=0.7,
    max_tokens=2000,
    streaming=True,
    keep_alive="5m",
)

from langchain.output_parsers import ResponseSchema, StructuredOutputParser
from langchain.schema import BaseOutputParser
from typing import Dict, Any
from datetime import datetime

# 1. 定义输出格式解析器
class TranslationOutputParser(BaseOutputParser):
    def parse(self, text: str) -> Dict[str, Any]:
        """解析模型输出的文本"""
        return {
            "translated_text": text.strip(),
            "length": len(text.strip()),
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        }

# 2. 创建格式化器
def format_translation_output(parsed_output: Dict[str, Any]) -> str:
    """格式化翻译输出"""
    return f"""
            翻译结果:
            -----------------
            译文: {parsed_output.get('translated_text', '')}
            字符数: {parsed_output.get('length', 0)}
            时间: {parsed_output.get('timestamp', '')}
            -----------------
            """

# 3. 更新提示模板
prompt = ChatPromptTemplate.from_template("""
            请将以下文本翻译成英文:
            {text}

            请只输出翻译结果，不要添加任何额外解释。
            """)

# 4. 创建带解析和格式化的链：LCEL表达式
chain = (
    prompt 
    | llm 
    | TranslationOutputParser()  # 解析响应
    | format_translation_output  # 格式化输出
)

# 5. 使用新的链
result = chain.invoke({"text": "我爱你"})
print(result)

print("end")